#include "sin_fit.h"

#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_multifit_nlin.h>

#include <stdlib.h>
#include <stdio.h>
#include <math.h>

#define N 256

void print_state(size_t iter, gsl_multifit_fdfsolver * s);

int main(void)
{
     const size_t n = N; // Number of data points.
     const size_t p = 3; // Number of fit parameters.

     // Create some data
     // ========================================================================
     // Setup random number genereator.
     gsl_rng_env_setup();
     const gsl_rng_type * type = gsl_rng_default;
     gsl_rng * r = gsl_rng_alloc(type);

     // Data to be fitted.
     double y[N];
     double sigma[N];
     double p_a = 4.0;
     double p_b = 0.8;
     double p_c = 0.0;
     double freq = 2.0;
     // TODO: NEED TO NORMALISE FREQ OUT OF THE X INDEX SOMEHOW...
     for (unsigned int i = 0; i < n; ++i)
     {
          double t = (double)i * freq / (double)n * 2 * M_PI;
          y[i] = p_a * sin(t + p_c) + p_c + gsl_ran_gaussian(r, 1.05);
          sigma[i] = 1.05;
          printf("data: %u %g %g\n", i, y[i], sigma[i]);
     }


     // Setup the solver.
     // ========================================================================
     // Set the solver type... ?!
     const gsl_multifit_fdfsolver_type * T = gsl_multifit_fdfsolver_lmsder;

     // Pointer to a newly allocated instance of a solver for type T,
     // n observations and p parameters.
     gsl_multifit_fdfsolver * s = gsl_multifit_fdfsolver_alloc(T, n, p);

     // Init solver s to use the fuction and derivative fdf and inital guess x.
     gsl_multifit_function_fdf f;
     f.f = &sin_f;
     f.df = &sin_df;
     f.fdf = &sin_fdf;
     f.n = n;
     f.p = p;
     struct Data data = { n, y, sigma };
     f.params = &data;
     double x_init[3] = { 3.0, 0.5, 0.0 };
     gsl_vector_view x = gsl_vector_view_array(x_init, p);
     gsl_multifit_fdfsolver_set(s, &f, &x.vector);


     // Solve
     // ========================================================================
     int status;
     unsigned int iter = 0;
     print_state(iter, s);

     do
     {
          iter++;
          status = gsl_multifit_fdfsolver_iterate(s);

          printf("status = %s\n", gsl_strerror(status));
          print_state(iter, s);

          if (status)
          {
               break;
          }

          status = gsl_multifit_test_delta(s->dx, s->x, 1e-4, 1e-4);
     }
     while (status == GSL_CONTINUE && iter < 500);
     printf("COMPLETE: status = %s\n", gsl_strerror(status));


     // Print the solution.
     // ========================================================================
     gsl_matrix * covar = gsl_matrix_alloc(p, p);
     gsl_multifit_covar(s->J, 0.0, covar);

     #define FIT(i) gsl_vector_get(s->x, i)
     #define ERR(i) sqrt(gsl_matrix_get(covar, i, i))
     double f_a = FIT(0);
     double f_b = FIT(1);
     double f_c = FIT(2);
     {
          double chi = gsl_blas_dnrm2(s->f);
          double dof = n - p;
          double c = GSL_MAX_DBL(1.0, chi / sqrt(dof));
          printf("=========================================================\n");
          printf("y = a * sin(freq * t + b) + c\n");
          printf("a = %f, b = %f, c = %f\n", p_a, p_b, p_c);
          printf("\n");
          printf("chisq per degree of freedom = %g\n", pow(chi, 2.0) / dof);
          printf("  a       = % .5f +/- % .5f\n", f_a, c * ERR(0));
          printf("  b       = % .5f +/- % .5f\n", f_b, c * ERR(1));
          printf("  c       = % .5f +/- % .5f\n", f_c, c * ERR(2));
          printf("=========================================================\n");
     }

     printf("FINISH WITH status = %s\n", gsl_strerror(status));

     gsl_multifit_fdfsolver_free (s);
     gsl_matrix_free (covar);
     gsl_rng_free (r);


     // Write data and fit to file.
     // ========================================================================
     FILE * data_file;
     data_file = fopen("data.txt", "w");
     for (unsigned int i = 0; i < n; ++i)
     {
          double t = (double)i * freq * 2 * M_PI / (double)n;
          double fit = f_a * sin(t + f_b) + f_c;
          fprintf(data_file, "%4u % .5g % .5g % .5g\n", i, y[i], sigma[i], fit);
     }
     fclose(data_file);

     return 0;
}



void print_state(size_t iter, gsl_multifit_fdfsolver * s)
{
     printf("iter: %3u x = % 15.8f % 15.8f % 15.8f "
            "|f(x)| = %f\n",
            (unsigned int)iter,
            gsl_vector_get(s->x, 0),
            gsl_vector_get(s->x, 1),
            gsl_vector_get(s->x, 2),
            gsl_blas_dnrm2(s->f));
}
